package com.catmiao.spark.acc

import org.apache.spark.util.{AccumulatorV2, LongAccumulator}
import org.apache.spark.{SparkConf, SparkContext}

import scala.collection.mutable

/**
 * @title: Spark01_RDD_Dep
 * @projectName spark_study
 * @description: 自定义累加器，word_count
 * @author ChengMiao
 * @date 2024/2/28 23:20
 */
object Spark01_Acc_02 {


  def main(args: Array[String]): Unit = {
    val sparkCon = new SparkConf().setMaster("local[*]").setAppName("acc")

    val sparkContext = new SparkContext(sparkCon)

    val rdd = sparkContext.makeRDD(List(
      "Hello World",
      "Hello Spark"
    ))

    // 创建累加器对象
    val wcAcc = new MyAccumulator

    // 像Spark进行注册
    sparkContext.register(wcAcc,"wordCountAcc")


    rdd.flatMap(_.split(" ")).foreach(
      word => {
        wcAcc.add(word)
      }
    )

    // 打印累加结果
    println(wcAcc.value)


    sparkContext.stop()
  }

  /**
   * 自定义数据累加器 wordCount
   * 1. 继承 AccumulatorV2 ,定义泛型
   *    - IN：累加器输入的数据类型
   *    - OUT：累加器返回的数据类型
   * 2. 重写方法
   */
  class MyAccumulator extends AccumulatorV2[String, mutable.Map[String, Long]] {

    private var wcMap = mutable.Map[String,Long]()

    // 判断是否为初始状态
    override def isZero: Boolean = {
      wcMap.isEmpty
    }

    // 复制累加器
    override def copy(): AccumulatorV2[String, mutable.Map[String, Long]] = {
     new MyAccumulator
    }

    // 重置累加器
    override def reset(): Unit = {
      wcMap.clear()
    }

    // 获取累加器需要计算的值
    override def add(word: String): Unit = {

      val newCount = wcMap.getOrElse(word,0L) + 1
      wcMap.update(word,newCount) // 累加数据
    }

    // 合并多个累加器 在Driver端
    override def merge(other: AccumulatorV2[String, mutable.Map[String, Long]]): Unit = {

      val map1 = wcMap
      val map2 = other.value
      map2.foreach{
        case (word,count) => {
          val newCount = wcMap.getOrElse(word, 0L) + count
          map1.update(word, newCount) // 累加数据
        }
      }
    }

    override def value: mutable.Map[String, Long] = {
      wcMap
    }
  }
}
